#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
RedFire 数据库迁移管理器
支持PostgreSQL和MySQL的数据库版本迁移
"""

import os
import re
import json
import logging
import asyncio
from datetime import datetime
from pathlib import Path
from typing import List, Dict, Any, Optional, Tuple
from dataclasses import dataclass

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


@dataclass
class Migration:
    """迁移文件信息"""
    version: str
    description: str
    file_path: str
    applied_at: Optional[datetime] = None
    checksum: Optional[str] = None


class MigrationManager:
    """数据库迁移管理器"""
    
    def __init__(self, db_config: Dict[str, Any], migrations_dir: str = None):
        self.db_config = db_config
        self.db_type = db_config.get('type', 'postgresql')
        self.migrations_dir = migrations_dir or os.path.join(
            os.path.dirname(__file__), 'versions'
        )
        self.migrations_table = 'schema_migrations'
        self.connection = None
    
    async def connect(self):
        """建立数据库连接"""
        if self.db_type == 'postgresql':
            import asyncpg
            self.connection = await asyncpg.connect(
                host=self.db_config['host'],
                port=self.db_config['port'],
                database=self.db_config['database'],
                user=self.db_config['username'],
                password=self.db_config['password']
            )
        elif self.db_type == 'mysql':
            import aiomysql
            self.connection = await aiomysql.connect(
                host=self.db_config['host'],
                port=self.db_config['port'],
                db=self.db_config['database'],
                user=self.db_config['username'],
                password=self.db_config['password'],
                charset='utf8mb4'
            )
        else:
            raise ValueError(f"不支持的数据库类型: {self.db_type}")
    
    async def close(self):
        """关闭数据库连接"""
        if self.connection:
            await self.connection.close()
    
    async def init_migrations_table(self):
        """初始化迁移记录表"""
        if self.db_type == 'postgresql':
            sql = f"""
            CREATE TABLE IF NOT EXISTS {self.migrations_table} (
                version VARCHAR(50) PRIMARY KEY,
                description TEXT,
                applied_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                checksum VARCHAR(64),
                execution_time INTEGER DEFAULT 0
            )
            """
        else:  # MySQL
            sql = f"""
            CREATE TABLE IF NOT EXISTS {self.migrations_table} (
                version VARCHAR(50) PRIMARY KEY,
                description TEXT,
                applied_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                checksum VARCHAR(64),
                execution_time INTEGER DEFAULT 0
            ) ENGINE=InnoDB CHARSET=utf8mb4
            """
        
        await self.execute_sql(sql)
        logger.info(f"迁移记录表 {self.migrations_table} 已初始化")
    
    async def execute_sql(self, sql: str, params=None):
        """执行SQL语句"""
        if self.db_type == 'postgresql':
            if params:
                return await self.connection.execute(sql, *params)
            else:
                return await self.connection.execute(sql)
        else:  # MySQL
            async with self.connection.cursor() as cursor:
                await cursor.execute(sql, params)
                return cursor.rowcount
    
    async def fetch_all(self, sql: str, params=None):
        """查询所有结果"""
        if self.db_type == 'postgresql':
            if params:
                return await self.connection.fetch(sql, *params)
            else:
                return await self.connection.fetch(sql)
        else:  # MySQL
            async with self.connection.cursor() as cursor:
                await cursor.execute(sql, params)
                return await cursor.fetchall()
    
    def get_migration_files(self) -> List[Migration]:
        """获取所有迁移文件"""
        migrations = []
        migrations_path = Path(self.migrations_dir)
        
        if not migrations_path.exists():
            logger.warning(f"迁移目录不存在: {self.migrations_dir}")
            return migrations
        
        # 匹配格式: V{version}__{description}.sql
        pattern = re.compile(r'^V(\d+(?:\.\d+)*)__(.+)\.sql$')
        
        for file_path in migrations_path.glob('*.sql'):
            match = pattern.match(file_path.name)
            if match:
                version = match.group(1)
                description = match.group(2).replace('_', ' ')
                
                migration = Migration(
                    version=version,
                    description=description,
                    file_path=str(file_path),
                    checksum=self.calculate_checksum(file_path)
                )
                migrations.append(migration)
        
        # 按版本号排序
        migrations.sort(key=lambda m: [int(x) for x in m.version.split('.')])
        return migrations
    
    def calculate_checksum(self, file_path: Path) -> str:
        """计算文件校验和"""
        import hashlib
        with open(file_path, 'rb') as f:
            return hashlib.md5(f.read()).hexdigest()
    
    async def get_applied_migrations(self) -> List[Migration]:
        """获取已应用的迁移"""
        sql = f"SELECT version, description, applied_at, checksum FROM {self.migrations_table} ORDER BY version"
        
        try:
            rows = await self.fetch_all(sql)
            migrations = []
            
            for row in rows:
                if self.db_type == 'postgresql':
                    migration = Migration(
                        version=row['version'],
                        description=row['description'],
                        file_path='',
                        applied_at=row['applied_at'],
                        checksum=row['checksum']
                    )
                else:  # MySQL
                    migration = Migration(
                        version=row[0],
                        description=row[1],
                        file_path='',
                        applied_at=row[2],
                        checksum=row[3]
                    )
                migrations.append(migration)
            
            return migrations
            
        except Exception as e:
            logger.warning(f"读取已应用迁移失败: {e}")
            return []
    
    async def apply_migration(self, migration: Migration) -> bool:
        """应用单个迁移"""
        logger.info(f"应用迁移: {migration.version} - {migration.description}")
        
        try:
            # 读取迁移文件
            with open(migration.file_path, 'r', encoding='utf-8') as f:
                sql_content = f.read()
            
            # 分割SQL语句（以分号分割，但忽略存储过程中的分号）
            statements = self.split_sql_statements(sql_content)
            
            start_time = datetime.now()
            
            if self.db_type == 'postgresql':
                # PostgreSQL 使用事务
                async with self.connection.transaction():
                    for statement in statements:
                        if statement.strip():
                            await self.connection.execute(statement)
            else:  # MySQL
                # MySQL 手动事务管理
                await self.execute_sql("START TRANSACTION")
                try:
                    for statement in statements:
                        if statement.strip():
                            await self.execute_sql(statement)
                    await self.execute_sql("COMMIT")
                except Exception:
                    await self.execute_sql("ROLLBACK")
                    raise
            
            execution_time = int((datetime.now() - start_time).total_seconds() * 1000)
            
            # 记录迁移历史
            await self.record_migration(migration, execution_time)
            
            logger.info(f"迁移 {migration.version} 应用成功，耗时 {execution_time}ms")
            return True
            
        except Exception as e:
            logger.error(f"迁移 {migration.version} 应用失败: {e}")
            return False
    
    def split_sql_statements(self, sql_content: str) -> List[str]:
        """分割SQL语句"""
        # 简单的SQL分割，可能需要更复杂的解析器来处理存储过程
        statements = []
        current_statement = ""
        in_delimiter = False
        
        for line in sql_content.split('\n'):
            line = line.strip()
            
            # 跳过注释和空行
            if not line or line.startswith('--') or line.startswith('#'):
                continue
            
            # 处理DELIMITER语句
            if line.upper().startswith('DELIMITER'):
                in_delimiter = not in_delimiter
                continue
            
            current_statement += line + '\n'
            
            # 如果在自定义分隔符中，查找分隔符
            if in_delimiter:
                if line.rstrip().endswith('//'):  # 假设使用//作为分隔符
                    statements.append(current_statement.rstrip('//\n'))
                    current_statement = ""
            else:
                # 普通分号分割
                if line.rstrip().endswith(';'):
                    statements.append(current_statement.rstrip(';\n'))
                    current_statement = ""
        
        # 添加最后一个语句
        if current_statement.strip():
            statements.append(current_statement.strip())
        
        return statements
    
    async def record_migration(self, migration: Migration, execution_time: int):
        """记录迁移历史"""
        sql = f"""
        INSERT INTO {self.migrations_table} 
        (version, description, checksum, execution_time)
        VALUES (%s, %s, %s, %s)
        """
        
        if self.db_type == 'postgresql':
            sql = sql.replace('%s', '$1')
            await self.connection.execute(
                sql, migration.version, migration.description, 
                migration.checksum, execution_time
            )
        else:
            await self.execute_sql(sql, (
                migration.version, migration.description,
                migration.checksum, execution_time
            ))
    
    async def get_pending_migrations(self) -> List[Migration]:
        """获取待应用的迁移"""
        all_migrations = self.get_migration_files()
        applied_migrations = await self.get_applied_migrations()
        applied_versions = {m.version for m in applied_migrations}
        
        pending = [m for m in all_migrations if m.version not in applied_versions]
        return pending
    
    async def migrate(self, target_version: str = None) -> bool:
        """执行迁移"""
        await self.connect()
        
        try:
            await self.init_migrations_table()
            
            pending_migrations = await self.get_pending_migrations()
            
            if not pending_migrations:
                logger.info("没有待应用的迁移")
                return True
            
            # 过滤到目标版本
            if target_version:
                pending_migrations = [
                    m for m in pending_migrations
                    if self.compare_versions(m.version, target_version) <= 0
                ]
            
            success_count = 0
            for migration in pending_migrations:
                if await self.apply_migration(migration):
                    success_count += 1
                else:
                    logger.error(f"迁移在版本 {migration.version} 处停止")
                    break
            
            logger.info(f"迁移完成，成功应用 {success_count}/{len(pending_migrations)} 个迁移")
            return success_count == len(pending_migrations)
            
        finally:
            await self.close()
    
    def compare_versions(self, version1: str, version2: str) -> int:
        """比较版本号"""
        v1_parts = [int(x) for x in version1.split('.')]
        v2_parts = [int(x) for x in version2.split('.')]
        
        # 补齐长度
        max_len = max(len(v1_parts), len(v2_parts))
        v1_parts.extend([0] * (max_len - len(v1_parts)))
        v2_parts.extend([0] * (max_len - len(v2_parts)))
        
        for v1, v2 in zip(v1_parts, v2_parts):
            if v1 < v2:
                return -1
            elif v1 > v2:
                return 1
        return 0
    
    async def rollback(self, target_version: str) -> bool:
        """回滚到指定版本（简单实现，仅删除记录）"""
        logger.warning("回滚功能需要手动实现回滚脚本")
        await self.connect()
        
        try:
            sql = f"DELETE FROM {self.migrations_table} WHERE version > %s"
            if self.db_type == 'postgresql':
                sql = sql.replace('%s', '$1')
                await self.connection.execute(sql, target_version)
            else:
                await self.execute_sql(sql, (target_version,))
            
            logger.info(f"已将迁移记录回滚到版本 {target_version}")
            logger.warning("请手动执行相应的回滚SQL语句")
            return True
            
        except Exception as e:
            logger.error(f"回滚失败: {e}")
            return False
        finally:
            await self.close()
    
    async def status(self):
        """显示迁移状态"""
        await self.connect()
        
        try:
            all_migrations = self.get_migration_files()
            applied_migrations = await self.get_applied_migrations()
            applied_versions = {m.version: m for m in applied_migrations}
            
            print("数据库迁移状态:")
            print("=" * 80)
            print(f"{'版本':<10} {'状态':<8} {'描述':<30} {'应用时间':<20}")
            print("-" * 80)
            
            for migration in all_migrations:
                if migration.version in applied_versions:
                    applied = applied_versions[migration.version]
                    status = "已应用"
                    applied_time = applied.applied_at.strftime('%Y-%m-%d %H:%M:%S') if applied.applied_at else "未知"
                else:
                    status = "待应用"
                    applied_time = ""
                
                print(f"{migration.version:<10} {status:<8} {migration.description:<30} {applied_time:<20}")
            
            pending_count = len([m for m in all_migrations if m.version not in applied_versions])
            print("-" * 80)
            print(f"总计: {len(all_migrations)} 个迁移，{len(applied_migrations)} 个已应用，{pending_count} 个待应用")
            
        finally:
            await self.close()


def create_migration(description: str, migrations_dir: str = None) -> str:
    """创建新的迁移文件"""
    if not migrations_dir:
        migrations_dir = os.path.join(os.path.dirname(__file__), 'versions')
    
    # 创建目录
    os.makedirs(migrations_dir, exist_ok=True)
    
    # 生成版本号（基于时间戳）
    timestamp = datetime.now().strftime('%Y%m%d%H%M%S')
    
    # 生成文件名
    safe_description = re.sub(r'[^\w\s-]', '', description).strip()
    safe_description = re.sub(r'[-\s]+', '_', safe_description)
    filename = f"V{timestamp}__{safe_description}.sql"
    
    file_path = os.path.join(migrations_dir, filename)
    
    # 创建迁移文件模板
    template = f"""-- RedFire 数据库迁移
-- 版本: {timestamp}
-- 描述: {description}
-- 创建时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}

-- ===========================================
-- 向前迁移 (Forward Migration)
-- ===========================================

-- 请在此处添加您的SQL语句
-- 例如:
-- CREATE TABLE example_table (
--     id SERIAL PRIMARY KEY,
--     name VARCHAR(100) NOT NULL,
--     created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
-- );

-- ===========================================
-- 回滚说明 (Rollback Instructions)
-- ===========================================

-- 如需回滚，请手动执行以下操作:
-- 1. 删除相关表: DROP TABLE IF EXISTS example_table;
-- 2. 删除相关索引: DROP INDEX IF EXISTS idx_example_name;
-- 3. 删除相关函数/存储过程
-- 4. 从迁移记录中删除此版本:
--    DELETE FROM schema_migrations WHERE version = '{timestamp}';

-- 注意: 生产环境中的回滚操作需要格外谨慎，建议先备份数据
"""
    
    with open(file_path, 'w', encoding='utf-8') as f:
        f.write(template)
    
    logger.info(f"迁移文件已创建: {file_path}")
    return file_path


# CLI工具
async def main():
    """命令行工具入口"""
    import sys
    import argparse
    
    parser = argparse.ArgumentParser(description='RedFire 数据库迁移工具')
    parser.add_argument('command', choices=['migrate', 'rollback', 'status', 'create'],
                       help='执行的命令')
    parser.add_argument('--version', '-v', help='目标版本号')
    parser.add_argument('--description', '-d', help='迁移描述（创建时使用）')
    parser.add_argument('--config', '-c', help='数据库配置文件路径')
    
    args = parser.parse_args()
    
    # 默认配置
    default_config = {
        'type': 'postgresql',
        'host': 'localhost',
        'port': 5432,
        'database': 'redfire_stage2',
        'username': 'redfire_prod',
        'password': 'your_password'
    }
    
    # 加载配置
    if args.config:
        import json
        with open(args.config, 'r') as f:
            db_config = json.load(f)
    else:
        db_config = default_config
    
    # 执行命令
    if args.command == 'create':
        if not args.description:
            print("错误: 创建迁移需要提供描述 (--description)")
            sys.exit(1)
        create_migration(args.description)
    
    else:
        manager = MigrationManager(db_config)
        
        if args.command == 'migrate':
            await manager.migrate(args.version)
        elif args.command == 'rollback':
            if not args.version:
                print("错误: 回滚需要提供目标版本 (--version)")
                sys.exit(1)
            await manager.rollback(args.version)
        elif args.command == 'status':
            await manager.status()


if __name__ == "__main__":
    asyncio.run(main())
